import copy
from typing import List,Dict

from .strategy_generator import StrategyGenerator,FollowingStrategyGenerator
from ..shard.placement_types import DTensorSpec,Placement,Replicate,PlacementStrategy,Shard,ShardingStrategy

__all__ = ['GetattrGenerator']


class GetitemGenerator(FollowingStrategyGenerator):
    """
    PlaceholderGenerator is a generic class to generate strategies for placeholder node.
    """
    def collate_strategies(self) -> List[ShardingStrategy]:
        strategy_list = []
        # 复制前面的策略,一输入一输出
        # 无法得到前面node的策略
        for index, strategy in enumerate(self.pre_strategies_vector):
            # 复制前一个策略,不要想复杂，一输入一输出,但是shard维度会变
            input_specs = copy.deepcopy(strategy.sharding_specs.output_specs)
            # 改变tensor_meta
            input_specs.tensor_meta = self.op_data['input_0'].data
            output_specs = DTensorSpec(mesh = self.device_mesh, placements=input_specs.placements,tensor_meta=self.op_data['output'].data)
            # output_specs不是list，所以要加[]
            sharding_specs = PlacementStrategy(input_specs=[input_specs],output_specs=output_specs)
            name = str(sharding_specs)
            strategy_list.append(ShardingStrategy(name=name,sharding_specs=sharding_specs,compute_cost =0))
        return strategy_list

